# Copyright 2020-2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""the module is used to process images."""

import copy

import numpy as np


def rand_init(a=0., b=1.):
    return np.random.rand() * (b - a) + a


def bbox_iou(bbox_a, bbox_b, offset=0):
    """Calculate Intersection-Over-Union(IOU) of two bounding boxes.

    Parameters
    ----------
    bbox_a : numpy.ndarray
        An ndarray with shape :math:`(N, 4)`.
    bbox_b : numpy.ndarray
        An ndarray with shape :math:`(M, 4)`.
    offset : float or int, default is 0
        The ``offset`` is used to control the whether the width(or height) is computed as
        (right - left + ``offset``).
        Note that the offset must be 0 for normalized bboxes, whose ranges are in ``[0, 1]``.

    Returns
    -------
    numpy.ndarray
        An ndarray with shape :math:`(N, M)` indicates IOU between each pairs of
        bounding boxes in `bbox_a` and `bbox_b`.

    """
    if bbox_a.shape[1] < 4 or bbox_b.shape[1] < 4:
        raise IndexError("Bounding boxes axis 1 must have at least length 4")

    tl = np.maximum(bbox_a[:, None, :2], bbox_b[:, :2])
    br = np.minimum(bbox_a[:, None, 2:4], bbox_b[:, 2:4])

    area_i = np.prod(br - tl + offset, axis=2) * (tl < br).all(axis=2)
    area_a = np.prod(bbox_a[:, 2:4] - bbox_a[:, :2] + offset, axis=1)
    area_b = np.prod(bbox_b[:, 2:4] - bbox_b[:, :2] + offset, axis=1)
    return area_i / (area_a[:, None] + area_b - area_i)


def is_iou_satisfied_constraint(box, crop_box, min_iou=None, max_iou=None):
    """Filter satisfied constraint IOU."""
    iou = bbox_iou(box, crop_box)
    return min_iou <= iou.min() and max_iou >= iou.max()


def choose_candidate_by_constraints(max_trial,
                                    input_w, input_h, image_w, image_h,
                                    jitter, box, use_constraints):
    """Choose candidate by constraints."""
    if use_constraints:
        constraints = ((0.1, None),
                       (0.3, None),
                       (0.5, None),
                       (0.7, None),
                       (0.9, None),
                       (None, 1),)
    else:
        constraints = ((None, None),)
    # add default candidate
    candidates = [(0, 0, input_w, input_h)]
    for constraint in constraints:
        min_iou, max_iou = constraint
        min_iou = -np.inf if min_iou is None else min_iou
        max_iou = np.inf if max_iou is None else max_iou

        for _ in range(max_trial):
            # box_data should have at least one box
            new_ar = float(input_w) / \
                     float(input_h) * rand_init(1 - jitter, 1 + jitter) / \
                     rand_init(1 - jitter, 1 + jitter)
            scale = rand_init(0.25, 2)

            if new_ar < 1:
                nh = int(scale * input_h)
                nw = int(nh * new_ar)
            else:
                nw = int(scale * input_w)
                nh = int(nw / new_ar)

            dx = int(rand_init(0, input_w - nw))
            dy = int(rand_init(0, input_h - nh))

            if box.size > 0:
                t_box = copy.deepcopy(box)
                t_box[:, [0, 2]] = t_box[:, [0, 2]] * float(nw) / float(
                    image_w) + dx
                t_box[:, [1, 3]] = t_box[:, [1, 3]] * float(nh) / float(
                    image_h) + dy

                crop_box = np.array((0, 0, input_w, input_h))
                if not is_iou_satisfied_constraint(
                        t_box, crop_box[np.newaxis],
                        min_iou, max_iou
                ):
                    continue
                else:
                    candidates.append((dx, dy, nw, nh))
            else:
                raise Exception("Box size should be greater than 0, but get {}.".format(box.size))
    return candidates


def correct_bbox_by_candidates(candidates,
                               input_w, input_h, image_w, image_h,
                               flip, box, box_data, allow_outside_center):
    """Calculate correct boxes.
    Args:
        candidates (list): object boxes, such as [(210, 291, 352, 403), .....].
        input_w (int): input image width.
        input_h (int): input image high.
        image_w (int): original image width.
        image_h (int): original image high.
        flip (bool): image flip flag. Default(True)
        box (list): its size is [2,4], save object boxes.
        box_data (ndarray): input box data, default is array([[0,0,0,0],[0,0,0,0],....])
        allow_outside_center: default is True
    Examples:
        >>> box_data, candidate = correct_bbox_by_candidates(
        candidates, input_w, input_h, image_w, image_h,
        flip, box, box_data, allow_outside_center
        )
    """
    while candidates:
        if len(candidates) > 1:
            # ignore default candidate which do not crop
            candidate = candidates.pop(
                np.random.randint(1, len(candidates)))
        else:
            candidate = candidates.pop(
                np.random.randint(0, len(candidates)))
        dx, dy, nw, nh = candidate
        t_box = copy.deepcopy(box)
        t_box[:, [0, 2]] = t_box[:, [0, 2]] * float(nw) / float(image_w) + dx
        t_box[:, [1, 3]] = t_box[:, [1, 3]] * float(nh) / float(image_h) + dy
        if flip:
            t_box[:, [0, 2]] = input_w - t_box[:, [2, 0]]

        if allow_outside_center:
            pass
        else:
            t_box = t_box[
                np.logical_and((t_box[:, 0] + t_box[:, 2]) / 2. >= 0.,
                               (t_box[:, 1] + t_box[:, 3]) / 2. >= 0.)]
            t_box = t_box[
                np.logical_and((t_box[:, 0] + t_box[:, 2]) / 2. <= input_w,
                               (t_box[:, 1] + t_box[:, 3]) / 2. <= input_h)]

        # re-correct x, y for case x,y < 0 reset to zero, after dx and dy, some box can smaller than zero
        t_box[:, 0:2][t_box[:, 0:2] < 0] = 0
        # re-correct w,h not higher than input size
        t_box[:, 2][t_box[:, 2] > input_w] = input_w
        t_box[:, 3][t_box[:, 3] > input_h] = input_h
        box_w = t_box[:, 2] - t_box[:, 0]
        box_h = t_box[:, 3] - t_box[:, 1]
        # discard invalid box: w or h smaller than 1 pixel
        t_box = t_box[np.logical_and(box_w > 1, box_h > 1)]

        if t_box.shape[0] > 0:
            # break if number of find t_box
            box_data[: len(t_box)] = t_box
            return box_data, candidate
    raise Exception("All box candidates can not satisfied re-correct bbox")
